﻿#include "InferOpenVino.h"

#include <QString>
#include <QFileInfo>
#include <random>

//#include <yaml-cpp/yaml.h>
#include <QString>
#include <fstream>
#include <QApplication>

#pragma warning(disable:4264)


InferOpenVino::~InferOpenVino()
{
}

DetectionResults InferOpenVino::runInfer(cv::Mat& src)
{
    int64 tick1 = cv::getTickCount();
    detections_.clear();
    //setSrcImage(src);
    if (!suc_loadLabel_ || !suc_loadModel_)
    {
        //qInfo() << "Please load model and classes files";
        return detections_;
    }

    auto inputShape = input_port_.get_shape();
    float factor = 0;
    cv::Mat blob = pre_process(&src, 640, &factor);

    ov::Tensor input_tensor(input_port_.get_element_type(), input_port_.get_shape(), blob.ptr(0));
    infer_.set_input_tensor(input_tensor);

    int64 tick2 = cv::getTickCount();

    infer_.infer();   // 同步机制

    int64 tick3 = cv::getTickCount();
    //qInfo() << "infer session->Run() time: " << t;

    auto output = infer_.get_output_tensor(0);
    auto output_shape = output.get_shape();
    //auto stride = output.get_strides();
    float* output_data = output.data<float>();

    std::vector<DetResult> result;
    switch (taskType_)
    {
    case infer_cls:
        qInfo() << "暂不支持";
        break;
    case infer_detect:
        result = post_process_v8(output_data, factor, output_shape[2]);
        //result = post_process_v5(output_data, factor, output_shape[2]);
        //result = post_process_v5_async(&infer_, factor);
        break;
    case infer_segment:
    {
        auto output1 = infer_.get_output_tensor(1);
        auto output1_shape = output1.get_shape();
        //result = post_process_v5(output_data, factor, output_shape[2]);
        //result = post_process_seg(output, output1, output_data, factor, output_shape[2]);
    }
        break;
    case infer_trace:
        break;
    case infer_obb:
        break;
    case infer_pos:
        break;
    default:
        break;
    }
    

    for (size_t j = 0; j < result.size(); j++) {
        Detection dtn;
        dtn.box = result[j].bbox;
        dtn.class_id = result[j].lable;
        dtn.className = labels_[result[j].lable];
        dtn.confidence = result[j].conf;

        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_int_distribution<int> dis(100, 255);
        dtn.color = cv::Scalar(dis(gen),
            dis(gen),
            dis(gen));

        detections_.push_back(dtn);
    }

    //draw_bbox(src, result);

    int64 tickEnd = cv::getTickCount();
    double t1 = (tick2 - tick1) * 1000 / cv::getTickFrequency();
    double t = (tick3 - tick2) * 1000 / cv::getTickFrequency();
    double t2 = (tickEnd - tick3) * 1000 / cv::getTickFrequency();
    //qInfo() << "pre : " << t1 << " --- infer: " << t << " --- post: " << t2;

	return detections_;
}

bool InferOpenVino::loadModel(std::string modelPath)
{
    std::wstring modelPathLoad = std::wstring(modelPath.begin(), modelPath.end());
    if (modelPath[1] != ':')
    {
        std::wstring appPath = QApplication::applicationDirPath().toStdWString();
        modelPathLoad = appPath + std::wstring(modelPath.begin(), modelPath.end());
    }

    int64 start = cv::getTickCount();
    ov::Core core;
    auto model = core.read_model(modelPathLoad);  // 直接读取onnx格式model

    // 读取ir格式：xml文件与bin文件
    //std::string binPath = modelPath.substr(0, modelPath.length() - 3) + "bin";
    //auto model = core.read_model("d:\\yolov8nIR-fp16\\yolov8n.xml", "d:\\yolov8nIR-fp16\\yolov8n.bin");


    auto compiled_model = core.compile_model(model, "AUTO",
    ov::hint::performance_mode(ov::hint::PerformanceMode::LATENCY),
    ov::hint::allow_auto_batching(true));

//auto compiled_model = core.compile_model(model, "AUTO",
//    ov::hint::performance_mode(ov::hint::PerformanceMode::
//        THROUGHPUT));

    //auto compiled_model = core.compile_model(model, "GPU");
    infer_ = compiled_model.create_infer_request();
    infer_.get_input_tensor().set_shape(std::vector<size_t>{1, 3, 640, 640});
    input_port_ = compiled_model.input();
    
    //input_port.get_shape();
    //input_port.get_element_type();
    

    int64 end = cv::getTickCount();
    double exeTime = 1000.0 * (end - start) / cv::getTickFrequency();
    qInfo() << "load model time : " << exeTime <<"ms";

    suc_loadModel_ = true;

	return suc_loadModel_;
}

bool InferOpenVino::loadLabels(std::string labelFilePath)
{
    std::wstring labelPathLoad = std::wstring(labelFilePath.begin(), labelFilePath.end());
    if (labelFilePath[1] != ':')
    {
        std::wstring appPath = QApplication::applicationDirPath().toStdWString();
        labelPathLoad = appPath + std::wstring(labelFilePath.begin(), labelFilePath.end());
    }

    int pos = labelFilePath.find_last_of('.');
    std::string str = labelFilePath.substr(pos + 1, labelFilePath.length() - pos - 1);
    labels_.clear();
    if (str == "txt")
    {
        parseTextFile(labelPathLoad);
    }
    else if (str == "yaml")
    {
        //parseYamlFile(labelFilePath);
        qInfo() << "不支持该格式";
    }

    suc_loadLabel_ = true;
    return suc_loadLabel_;
}

void InferOpenVino::setSrcImage(cv::Mat& src)
{
    srcImage_ = src;
}

void InferOpenVino::draw_bbox(cv::Mat& img, std::vector<DetResult>& res)
{
    for (size_t j = 0; j < res.size(); j++) {
        cv::rectangle(img, res[j].bbox, cv::Scalar(255, 0, 255), 2);
        cv::putText(img, std::to_string(res[j].lable) + "-" + std::to_string(res[j].conf),
            cv::Point(res[j].bbox.x, res[j].bbox.y - 1), cv::FONT_HERSHEY_PLAIN,
            1.2, cv::Scalar(0, 0, 255), 2);
    }
}

cv::Mat InferOpenVino::pre_process(cv::Mat* img, int length, float* factor)
{
    //cv::Mat mat;
    int rh = img->rows;
    int rw = img->cols;
    int max_image_length = rw > rh ? rw : rh;
    cv::Mat max_image = cv::Mat::zeros(max_image_length, max_image_length, CV_8UC3);
    img->copyTo(cv::Mat(max_image, cv::Rect(0, 0, rw, rh)));    
    *factor = (float)((float)max_image_length / (float)length);
    auto inputShape = input_port_.get_shape();
    int64 s = cv::getTickCount();
    cv::Mat blob = cv::dnn::blobFromImage(max_image, 1.0 / 255.0, cv::Size(inputShape[2], inputShape[3]), cv::Scalar(), true);
    int64 e = cv::getTickCount();
    double t = (e - s) * 1000 / cv::getTickFrequency();
    //qInfo() << "blob: " << t;

    return blob;
}

std::vector<DetResult> InferOpenVino::post_process_v8(float* result, float factor, int outputLength)
{
    std::vector<cv::Rect> position_boxes;
    std::vector <int> class_ids;
    std::vector <float> confidences;
    // Preprocessing output results
    

    for (int i = 0; i < outputLength; i++)
    {
        for (int j = 4; j < 4 + labels_.size(); j++)
        {
            float source = result[outputLength * j + i];
            int label = j - 4;
            if (source > modelConfidenceThreshold_)
            {
                float maxSource = source;
                float cx = result[outputLength * 0 + i];
                float cy = result[outputLength * 1 + i];
                float ow = result[outputLength * 2 + i];
                float oh = result[outputLength * 3 + i];
                int x = (int)((cx - 0.5 * ow) * factor);
                int y = (int)((cy - 0.5 * oh) * factor);
                int width = (int)(ow * factor);
                int height = (int)(oh * factor);

                cv::Rect box(x, y, width, height);

                position_boxes.push_back(box);
                class_ids.push_back(label);
                confidences.push_back(maxSource);
            }
        }
    }

    std::vector<int> indices;
    cv::dnn::NMSBoxes(position_boxes, confidences, modelScoreThreshold_, modelNMSThreshold_, indices);
    std::vector<DetResult> re;
    for (int i = 0; i < indices.size(); i++)
    {
        int index = indices[i];
        DetResult det(position_boxes[index], confidences[index], class_ids[index]);
        re.push_back(det);
    }

    return re;
}

std::vector<DetResult> InferOpenVino::post_process_v5(float* result, float factor, int outputLength)
{
    std::vector<cv::Rect> boxes;
    std::vector <int> class_ids;
    std::vector<float> class_scores;
    std::vector <float> confidences;
    // Preprocessing output results

    auto output = infer_.get_output_tensor(0);
    auto output_shape = output.get_shape();
    // 25200 x 85 Matrix 
    cv::Mat output_buffer(output_shape[1], output_shape[2], CV_32F, output.data());
    for (int i = 0; i < output_buffer.rows; i++) {
        float confidence = output_buffer.at<float>(i, 4);
        if (confidence < modelConfidenceThreshold_) {
            continue;
        }
        cv::Mat classes_scores = output_buffer.row(i).colRange(output_shape[2] - labels_.size(), output_shape[2]);
        cv::Point class_id;
        double score;
        cv::minMaxLoc(classes_scores, NULL, &score, NULL, &class_id);

        // class score: 0~1
        if (score > modelConfidenceThreshold_)
        {
            float cx = output_buffer.at<float>(i, 0);
            float cy = output_buffer.at<float>(i, 1);
            float w = output_buffer.at<float>(i, 2);
            float h = output_buffer.at<float>(i, 3);
            int x = (int)((cx - 0.5 * w) * factor);
            int y = (int)((cy - 0.5 * h) * factor);
            int width = (int)(w * factor);
            int height = (int)(h * factor);

            cv::Rect box(x, y, width, height);

            boxes.push_back(box);
            class_ids.push_back(class_id.x);
            class_scores.push_back(score);
            confidences.push_back(confidence);
        }
    }

    std::vector<int> indices;
    cv::dnn::NMSBoxes(boxes, confidences, modelScoreThreshold_, modelNMSThreshold_, indices);

    std::vector<DetResult> re;
    for (auto index : indices)
    {
        DetResult det(boxes[index], confidences[index], class_ids[index]);
        re.push_back(det);
    }

    return re;
}

std::vector<DetResult> InferOpenVino::post_process_v5_async(ov::InferRequest* infer, float factor)
{
    std::vector<cv::Rect> boxes;
    std::vector <int> class_ids;
    std::vector<float> class_scores;
    std::vector <float> confidences;
    // Preprocessing output results

    auto output = infer->get_output_tensor(0);
    auto output_shape = output.get_shape();
    // 25200 x 85 Matrix 
    cv::Mat output_buffer(output_shape[1], output_shape[2], CV_32F, output.data());
    for (int i = 0; i < output_buffer.rows; i++) {
        float confidence = output_buffer.at<float>(i, 4);
        if (confidence < modelConfidenceThreshold_) {
            continue;
        }
        cv::Mat classes_scores = output_buffer.row(i).colRange(output_shape[2] - labels_.size(), output_shape[2]);
        cv::Point class_id;
        double score;
        cv::minMaxLoc(classes_scores, NULL, &score, NULL, &class_id);

        // class score: 0~1
        if (score > modelConfidenceThreshold_)
        {
            float cx = output_buffer.at<float>(i, 0);
            float cy = output_buffer.at<float>(i, 1);
            float w = output_buffer.at<float>(i, 2);
            float h = output_buffer.at<float>(i, 3);
            int x = (int)((cx - 0.5 * w) * factor);
            int y = (int)((cy - 0.5 * h) * factor);
            int width = (int)(w * factor);
            int height = (int)(h * factor);

            cv::Rect box(x, y, width, height);

            boxes.push_back(box);
            class_ids.push_back(class_id.x);
            class_scores.push_back(score);
            confidences.push_back(confidence);
        }
    }

    std::vector<int> indices;
    cv::dnn::NMSBoxes(boxes, confidences, modelScoreThreshold_, modelNMSThreshold_, indices);

    std::vector<DetResult> re;
    for (auto index : indices)
    {
        DetResult det(boxes[index], confidences[index], class_ids[index]);
        re.push_back(det);
    }

    return re;
}

using namespace cv;
std::vector<cv::Scalar> colors = { Scalar(255, 0, 0), Scalar(255, 0, 255), Scalar(170, 0, 255), Scalar(255, 0, 85),
                                   Scalar(255, 0, 170), Scalar(85, 255, 0), Scalar(255, 170, 0), Scalar(0, 255, 0),
                                   Scalar(255, 255, 0), Scalar(0, 255, 85), Scalar(170, 255, 0), Scalar(0, 85, 255),
                                   Scalar(0, 255, 170), Scalar(0, 0, 255), Scalar(0, 255, 255), Scalar(85, 0, 255) };

float sigmoid_function(float a) {
    float b = 1. / (1. + exp(-a));
    return b;
}
std::vector<DetResult> InferOpenVino::post_process_seg(ov::Tensor output0, ov::Tensor output1, float* result, float factor, int outputLength)
{
    auto output0_shape = output0.get_shape();
    auto output1_shape = output1.get_shape();

    // -------- 将输出进行转置：[116,8400]--->[8400，116] --------
    // 116=4+80+32 [x,y,w,h]+80类别+32Mask
    cv::Mat outputBuffer(output0_shape[1], output0_shape[2], CV_32F, output0.data<float>());
    cv::Mat proto(output1_shape[1], output1_shape[2] * output1_shape[3], CV_32F, output1.data<float>()); //[32,25600]
    cv::transpose(outputBuffer, outputBuffer); //[8400,116]

    std::vector<cv::Rect> position_boxes;
    std::vector <int> class_ids;
    std::vector <float> confidences;
    std::vector<cv::Mat> mask_confs;
    // Preprocessing output results

    for (int i = 0; i < outputBuffer.rows; i++)
    {
        /*cv::Mat rowData = outputBuffer.row(i).colRange(0, 116);
        cv::Mat classesScore = rowData.row(0).colRange(4,84);*/
        cv::Mat rowData = outputBuffer.row(i).colRange(0, outputBuffer.cols);
        cv::Mat classesScore = rowData.row(0).colRange(4, labels_.size() + 4);
        cv::Point class_id;
        double maxClassScore;
        minMaxLoc(classesScore, 0, &maxClassScore, 0, &class_id);
        if (maxClassScore > modelConfidenceThreshold_)
        {
            float maxSource = maxClassScore;
            float cx = rowData.at<float>(0,0);
            float cy = rowData.at<float>(0, 1);
            float ow = rowData.at<float>(0, 2);
            float oh = rowData.at<float>(0,3);
            int x = (int)((cx - 0.5 * ow) * factor);
            int y = (int)((cy - 0.5 * oh) * factor);
            int width = (int)(ow * factor);
            int height = (int)(oh * factor);

            cv::Rect box(x, y, width, height);

            position_boxes.push_back(box);
            class_ids.push_back(class_id.x);
            confidences.push_back(maxSource);

            // mask
            cv::Mat mask_conf = outputBuffer.row(i).colRange(84, 116);
            mask_confs.push_back(mask_conf);
        }
    }
    
    std::vector<int> indices;
    cv::dnn::NMSBoxes(position_boxes, confidences, modelScoreThreshold_, modelNMSThreshold_, indices);
    std::vector<DetResult> re;
    for (int i = 0; i < indices.size(); i++)
    {
        int index = indices[i];
        DetResult det(position_boxes[index], confidences[index], class_ids[index]);
        re.push_back(det);
    }

    //int64 ss = cv::getTickCount();
    cv::Mat rgb_mask = cv::Mat::zeros(srcImage_.size(), srcImage_.type());
    cv::Mat masked_img;
    cv::RNG rng;
    for (size_t i = 0; i < indices.size(); i++) {
        // Visualize the objects
        int index = indices[i];
        int class_id = class_ids[index];

        // Visualize the Masks
        Mat m = mask_confs[index] * proto;
        for (int col = 0; col < m.cols; col++) {
            m.at<float>(0, col) = sigmoid_function(m.at<float>(0, col));
        }
        cv::Mat m1 = m.reshape(1, 160); // 1x25600 -> 160x160
        int x1 = std::max(0, position_boxes[index].x);
        int y1 = std::max(0, position_boxes[index].y);
        int x2 = std::max(0, position_boxes[index].br().x);
        int y2 = std::max(0, position_boxes[index].br().y);
        int mx1 = int(x1 / factor * 0.25);
        int my1 = int(y1 / factor * 0.25);
        int mx2 = int(x2 / factor * 0.25);
        int my2 = int(y2 / factor * 0.25);

        cv::Mat mask_roi = m1(cv::Range(my1, my2), cv::Range(mx1, mx2));
        cv::Mat rm, det_mask;
        cv::resize(mask_roi, rm, cv::Size(x2 - x1, y2 - y1));

        for (int r = 0; r < rm.rows; r++) {
            for (int c = 0; c < rm.cols; c++) {
                float pv = rm.at<float>(r, c);
                if (pv > 0.5) {
                    rm.at<float>(r, c) = 1.0;
                }
                else {
                    rm.at<float>(r, c) = 0.0;
                }
            }
        }
        rm = rm * rng.uniform(0, 255);
        rm.convertTo(det_mask, CV_8UC1);
        if ((y1 + det_mask.rows) >= srcImage_.rows) {
            y2 = srcImage_.rows - 1;
        }
        if ((x1 + det_mask.cols) >= srcImage_.cols) {
            x2 = srcImage_.cols - 1;
        }

        cv::Mat mask = cv::Mat::zeros(cv::Size(srcImage_.cols, srcImage_.rows), CV_8UC1);
        det_mask(cv::Range(0, y2 - y1), cv::Range(0, x2 - x1)).copyTo(mask(cv::Range(y1, y2), cv::Range(x1, x2)));
        add(rgb_mask, cv::Scalar(rng.uniform(0, 255), rng.uniform(0, 255), rng.uniform(0, 255)), rgb_mask, mask);
        addWeighted(srcImage_, 0.5, rgb_mask, 0.5, 0, masked_img);
    }
    masked_img.copyTo(srcImage_);

    /*int64 ee = cv::getTickCount();
    double ttt = 1000.0*(ee - ss) / cv::getTickFrequency();
    qInfo() << "post process segment: " << ttt << "ms";*/

   
    return re;
}


void InferOpenVino::parseTextFile(std::wstring filePath)
{
    std::ifstream inputFile(filePath);
    if (inputFile.is_open())
    {
        std::string classLine;
        while (std::getline(inputFile, classLine))
        {
            int p = classLine.find(":");
            if (p < 0)
            {
                continue;
            }
            classLine = classLine.substr(p + 1, classLine.length() - 1);
            if (classLine.length() < 1)
            {
                continue;
            }
            labels_.push_back(classLine);
        }
        inputFile.close();
    }
}

//void InferOpenVino::parseYamlFile(std::string filePath)
//{
//    // 读取文件，并解析
//    YAML::Node classesData = YAML::LoadFile(filePath);
//
//    // 检索yaml中的值
//    const YAML::Node names = classesData["names"];
//    for (auto& item : names)
//    {
//        if (item.IsDefined())   // 自定义格式
//        {
//            std::string s = QString(item.as<std::string>().c_str()).toLocal8Bit().constData();
//            labels_.push_back(s);
//        }
//        else  // coco.yaml标准格式
//        {
//            std::string s = QString(item.second.as<std::string>().c_str()).toLocal8Bit().constData();
//            labels_.push_back(s);
//        }
//    }
//}

